from gurobipy import *
import numpy as np
from traffic_engineer import *
from traffic_engineering_class import *

class BoundedWCMPTrafficEngineerStrong(TrafficEngineer):
	## Takes in two new variables:
	## update_frequency - how much time in-between every traffic engineering events
	def __init__(self, aurora_network, all_interblock_paths, update_frequency, training_length, reduce_multihop=False, ignore_flows=None):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		self.flow_pairs_to_ignore = ignore_flows
		self.update_frequency = update_frequency
		self.training_length = training_length
		self.reduce_multihop = reduce_multihop
		return

	def get_filename_param(self):
		return "bwcmptestrong_r{}t{}".format(self.update_frequency, self.training_length)

	# returns a string description of this class
	def get_string(self):
		return "Bounded WCMP Traffic Engineer Strong"

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.ROBUST_TE

	def _minimize_mlu(self, logical_topology_adj_matrix, max_tm, ingress_bound, egress_bound):
		nblocks = self.aurora_network.get_num_blocks()
		
		# Step 1 : Initialize the optimization variables
		model = Model("Bounded WCMP Minimize MLU")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mlu")
		routing_weights_var = [None] * (nblocks + 1)
		for index in range(nblocks + 1):
			routing_weights_var[index] = model.addVar(lb=0., ub=1, obj=0, vtype=GRB.CONTINUOUS, name="w_{}".format(index))

		## Step 2 : Start programming the constraints
		## Constraint Type 1 : weights must sum to 1
		weight_constraint = LinExpr()
		for index in range(nblocks + 1):
			weight_constraint += routing_weights_var[index]
		model.addConstr(lhs=weight_constraint,sense=GRB.EQUAL,rhs=1)
		## Constraint Type 2 : mlu constraint
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					mlu_constraint = LinExpr()
					Zi = egress_bound[i]
					Zj = ingress_bound[j]
					Uij = max_tm[i][j]
					w0 = routing_weights_var[0]
					wi = routing_weights_var[i + 1]
					wj = routing_weights_var[j + 1]
					mlu_constraint = Zi * wj + Zj * wi + w0 * Uij
					link_capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=mlu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * link_capacity * logical_topology_adj_matrix[i][j])
		
		## Step 3 : insert the objective function
		model.setObjective(mlu, GRB.MINIMIZE)
		try:
			model.optimize()
			## try to minimize sensitivty when computing the routing weights
			print("Bounded WCMP - MLU : {}".format(mlu.x))
			print("Guaranteed direct hop weight is : {}".format(routing_weights_var[0].x))
			routing_weights = {}
			block_pair_utilization = np.zeros((nblocks, nblocks))		
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						direct_path = (i,j)
						direct_path_weight = routing_weights_var[0].x + routing_weights_var[i + 1].x + routing_weights_var[j + 1].x
						routing_weights[direct_path] = direct_path_weight
						for k in range(nblocks):
							if k != i and k != j:
								indirect_path = (i,k,j)
								routing_weights[indirect_path] = routing_weights_var[k + 1].x
			return routing_weights, mlu.x
		except GurobiError as e:
			print ("Bounded WCMP - Error code " + str(e. errno ) + ": " + str(e))
		except AttributeError :
			print ("Bounded WCMP - Encountered an attribute error ")
		return None

	## Given a specific level of MLU performance to preserve.
	def _minimize_maximum_sensitivity(self, logical_topology_adj_matrix, preserved_mlu, max_tm, ingress_bound, egress_bound):
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("Desensitize")
		model.setParam( 'OutputFlag', False )
		routing_weight_vars = [None] * (nblocks + 1) 
		beta = model.addVar(lb=0., ub=GRB.INFINITY, obj=1, vtype=GRB.CONTINUOUS, name="beta")
		## Step 1 : Start programming the constraints
		## Constraint Type 1 : weights must sum to 1
		weight_constraint = LinExpr()
		for index in range(nblocks + 1):
			routing_weight_vars[index] = model.addVar(lb=0., ub=1, obj=0, vtype=GRB.CONTINUOUS, name="w_{}".format(index))
			weight_constraint += routing_weight_vars[index]
		model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1)
		
		## Constraint Type 2 : mlu constraint
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					mlu_constraint = LinExpr()
					Zi = egress_bound[i]
					Zj = ingress_bound[j]
					Uij = max_tm[i][j]
					w0 = routing_weight_vars[0]
					wi = routing_weight_vars[i + 1]
					wj = routing_weight_vars[j + 1]
					mlu_constraint = Zi * wj + Zj * wi + w0 * Uij
					link_capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=mlu_constraint, sense=GRB.LESS_EQUAL, rhs=preserved_mlu * link_capacity * logical_topology_adj_matrix[i][j])
		
		## stage 3 : add the sensitivity constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					direct_path_capacity = logical_topology_adj_matrix[i][j]						
					model.addConstr(lhs=routing_weight_vars[0] + routing_weight_vars[i + 1] + routing_weight_vars[j + 1], sense=GRB.LESS_EQUAL, rhs=beta * direct_path_capacity)
					for k in range(nblocks):
						if k != i and k != j:
							path_capacity = min(logical_topology_adj_matrix[i][k], logical_topology_adj_matrix[k][j])
							model.addConstr(lhs=routing_weight_vars[k + 1], sense=GRB.LESS_EQUAL, rhs=beta * path_capacity)
		model.setObjective(beta, GRB.MINIMIZE)
		try: 
			model.optimize()
			print("Minimum achievable sensivity is : {}".format(beta.x))
			print("The routing weights distribution is : {}".format([entry.x for entry in routing_weight_vars]))
			routing_weights = {}
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						direct_path = (i,j)
						direct_path_weight = routing_weight_vars[0].x + routing_weight_vars[i + 1].x + routing_weight_vars[j + 1].x
						routing_weights[direct_path] = direct_path_weight
						for k in range(nblocks):
							if k != i and k != j:
								indirect_path = (i,k,j)
								routing_weights[indirect_path] = routing_weight_vars[k + 1].x
			return routing_weights, beta.x
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error Minimize maximum sensitivity")
			return None

	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrices):
		## newly added for robustness
		nblocks = self.aurora_network.get_num_blocks()
		egress_bound = [0] * nblocks
		ingress_bound = [0] * nblocks
		## Step 1 : compute the pod egress and ingress bounds
		max_tm = np.zeros((nblocks, nblocks,))
		for tm in traffic_matrices:
			for pod in range(nblocks):
				pod_egress = sum(tm[pod])
				pod_ingress = 0
				for target_pod in range(nblocks):
					if target_pod != pod:
						pod_ingress += tm[target_pod][pod]
						max_tm[pod][target_pod] = max(max_tm[pod][target_pod], tm[pod][target_pod])
				egress_bound[pod] = max(egress_bound[pod], pod_egress)
				ingress_bound[pod] = max(ingress_bound[pod], pod_ingress)

		## Step 2 : Compute the worst case MLU
		routing_weights, mlu = self._minimize_mlu(logical_topology_adj_matrix, max_tm, ingress_bound, egress_bound)
		routing_weights, max_sensitivity = self._minimize_maximum_sensitivity(logical_topology_adj_matrix, 1.05 * mlu, max_tm, ingress_bound, egress_bound)
		return routing_weights